import torch
import numpy as np
import random

def init_seeds(seed=0):
    np.random.seed(seed)
    random.seed(seed)
    torch.manual_seed(seed) # cpu
    torch.cuda.manual_seed_all(seed)  # gpu
    torch.backends.cudnn.deterministic = True  # consistent results on the cpu and gpu

    # 设置随机种子  rank = -1
    # 在神经网络中，参数默认是进行随机初始化的。如果不设置的话每次训练时的初始化都是随机的，
    # 导致结果不确定。如果设置初始化，则每次初始化都是固定的
    torch.manual_seed(seed)

    torch.backends.cudnn.benchmark = False

def cuda2cpu(pred):
    """cuda 变量 转cpu"""
    if type(pred) == list:
        return pred
    if pred.is_cuda:
        pred_cpu = pred.cpu().numpy()
    else:
        pred_cpu = pred.numpy()
    return pred_cpu

def cv2ToTorch(img):
    """cv2的格式转torch 的tensor"""
    img = img[:, :, ::-1].transpose(2, 0, 1)  # BGR to RGB, to cxhxw
    img = np.ascontiguousarray(img, dtype=np.float32)
    img = torch.from_numpy(img)
    return img

def select_device(device=''):
    return torch.device('cuda:{}'.format(device) if torch.cuda.is_available() else 'cpu')